from driving_gym.environment.agent.done.base_done import BaseDone
from driving_gym.environment.scenario.carla_scenario import CarlaScenario
from driving_gym.misc.roach.core.task_actor.ego_vehicle.ego_vehicle_handler import (
    EgoVehicleHandler,
    TaskVehicle,
)
from driving_gym.misc.roach.core.task_actor.ego_vehicle.terminal.valeo_no_det_px import (
    ValeoNoDetPx,
)
from driving_gym.misc.roach.utils.traffic_light import TrafficLightHandler
from driving_gym.simulation.adapter.carla.carla_adapter import CarlaAdapter
from driving_gym.simulation.adapter.carla.carla_config import (
    CarlaActorConfig,
    carla_util,
)


class CarlaRoachDone(BaseDone):
    adapter: CarlaAdapter

    def __init__(self, config: dict, adapter=None):
        super().__init__(config, adapter)

        self.terminal_handler = None
        self._tick_ego = False

    def reset(self):
        self.terminal_handler = None
        if TrafficLightHandler.carla_map is None:
            TrafficLightHandler.reset(self.adapter.world)
        TaskVehicle.VEHICLES[self.actor_id] = None

        snap_shot = self.adapter.world.get_snapshot()
        self._timestamp = {
            "step": 0,
            "frame": snap_shot.timestamp.frame,
            "relative_wall_time": 0.0,
            "wall_time": snap_shot.timestamp.platform_timestamp,
            "relative_simulation_time": 0.0,
            "simulation_time": snap_shot.timestamp.elapsed_seconds,
            "start_frame": snap_shot.timestamp.frame,
            "start_wall_time": snap_shot.timestamp.platform_timestamp,
            "start_simulation_time": snap_shot.timestamp.elapsed_seconds,
        }
        return super().reset()

    def get_done(self, snapshot, scenario: CarlaScenario):
        if self.terminal_handler is None:
            self.terminal_handler = self._create_handler(scenario)

        self._update_timestamp()
        if self._tick_ego:
            self.terminal_handler._ego_vehicle.tick(self._timestamp)

        done, timeout, _, info = self.terminal_handler.get(self._timestamp)

        reason = "NOT_DONE"
        if info["blocked"]:
            reason = "BLOCKED"
        elif info["offroad"]:
            reason = "OFFROAD"
        elif info["run_red_light"]:
            reason = "RUN_RED_LIGHT"
        elif info["collision"]:
            reason = "COLLISION"
        elif info["run_stop_sign"]:
            reason = "RUN_STOP_SIGN"
        elif info["collision_px"]:
            reason = "COLLISION_PX"
        elif info["timeout"]:
            reason = "TIMEOUT"

        return done, reason

    def _create_handler(self, scenario: CarlaScenario):
        ego_vehicle = self.adapter.get_actor(self.actor_id)
        carla_map = self.adapter._map

        task_vehicle = TaskVehicle.VEHICLES.get(self.actor_id, None)
        if task_vehicle is None:
            spawn_transforms = EgoVehicleHandler._get_spawn_points(carla_map)
            task = scenario.get_task(self.actor_id)
            route = task.get("route", None) or [task["end"]]
            target_transforms = [
                carla_util.carla_tranform(
                    CarlaActorConfig.parse_location(loc, carla_map)
                )
                for loc in route
            ]

            task_vehicle = TaskVehicle(
                ego_vehicle, target_transforms, spawn_transforms, scenario.endless
            )
            TaskVehicle.VEHICLES[self.actor_id] = task_vehicle
            self._tick_ego = True

        exploration_suggest = self.config.get("exploration_suggest", False)
        eval_mode = self.config.get("eval_mode", False)
        terminal_handler = ValeoNoDetPx(task_vehicle, exploration_suggest, eval_mode)
        return terminal_handler

    def _update_timestamp(self):
        snap_shot = self.adapter.world.get_snapshot()
        self._timestamp["step"] = (
            snap_shot.timestamp.frame - self._timestamp["start_frame"]
        )
        self._timestamp["frame"] = snap_shot.timestamp.frame
        self._timestamp["wall_time"] = snap_shot.timestamp.platform_timestamp
        self._timestamp["relative_wall_time"] = (
            self._timestamp["wall_time"] - self._timestamp["start_wall_time"]
        )
        self._timestamp["simulation_time"] = snap_shot.timestamp.elapsed_seconds
        self._timestamp["relative_simulation_time"] = (
            self._timestamp["simulation_time"]
            - self._timestamp["start_simulation_time"]
        )
